package cn.itcast.flink.base.test;

import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.util.Iterator;

/**
 * Author itcast
 * Date 2021/7/28 19:01
 * desc：
 * 主要是创建一个 checkpoint ，将报错的offset保存到checkpoint中
 * 提供每5秒钟报一次错误
 */
public class OperateStateDemo {
    public static void main(String[] args) throws Exception {
        //1.创建流环境，便于观察设置并行度为 1
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        //2.开启checkpoint ，并将状态保存到 file:///D:/chk
        env.enableCheckpointing(1000);
        env.setStateBackend(new FsStateBackend("file:///d:/chk"));
        //3.设置checkpoint的配置 外部chk，仅一次语义等
        //env.getCheckpointConfig().setCheckpointTimeout(60000);
        //env.getCheckpointConfig().enableExternalizedCheckpoints(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
        //env.getCheckpointConfig().setMaxConcurrentCheckpoints(1);
        //env.getCheckpointConfig().setTolerableCheckpointFailureNumber(1);
        //4.开启重启策略 3秒钟尝试重启3次
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3, 3000));
        //5.添加数据源比如 MyMonitorKafkaSource , 实例化创建 MyMonitorKafkaSource
        DataStreamSource<String> source = env.addSource(new MySourceFunction());
        //6.打印输出
        source.print();
        //7.执行
        env.execute();

        //创建 MyMonitorKafkaSource 继承 RichParallelSourceFunction<String> 并实现 CheckpointedFunction
        //重写initializeState方法 ListStateDescriptor 状态描述和通过context获取 offsetState
        //重写run方法 读取出 offset 并 循环读取offset+=1，拿到执行的核心编号，输出(核编号和offset),一秒一条,每5条模拟一个异常
        //重写cancel方法
        //重写snapshotState方法 ， 清空 offsetState ，并将最新的offset添加进去
    }

    private static class MySourceFunction extends RichParallelSourceFunction<String> implements CheckpointedFunction {
        ListState<Long> listState = null;
        //设置 flag
        boolean flag = true;
        //当前的 offset
        Long currentOffset = 0L;

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            ListStateDescriptor<Long> offsetDesc = new ListStateDescriptor<>("offset", Long.class);
            listState = context.getOperatorStateStore().getListState(offsetDesc);
        }

        @Override
        public void run(SourceContext<String> ctx) throws Exception {
            //遍历出liststate内容
            Iterator<Long> iterator = listState.get().iterator();
            //如果有offset就赋值给当前 offset
            if (iterator.hasNext()) {
                currentOffset = iterator.next();
            }
            //
            while (flag) {
                currentOffset += 1L;
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                String output = "indexOfThisSubtask:" + indexOfThisSubtask + " offset:" + currentOffset;
                Thread.sleep(2000);
                ctx.collect(output);
                if (currentOffset % 5 == 0) {
                    System.out.println("程序遇到异常了......");
                    throw new Exception("此处遇到一个bug！");
                }
            }

        }

        @Override
        public void cancel() {
            flag = false;
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            listState.clear();
            listState.add(currentOffset);
        }
    }
}
